///
/// Copyright (C) 2018, Cyberhaven
///
/// Permission is hereby granted, free of charge, to any person obtaining a copy
/// of this software and associated documentation files (the "Software"), to deal
/// in the Software without restriction, including without limitation the rights
/// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
/// copies of the Software, and to permit persons to whom the Software is
/// furnished to do so, subject to the following conditions:
///
/// The above copyright notice and this permission notice shall be included in all
/// copies or substantial portions of the Software.
///
/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
/// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
/// SOFTWARE.
///

#include <boost/algorithm/string/predicate.hpp>

#include <s2e/ConfigFile.h>
#include <s2e/S2E.h>
#include <s2e/Utils.h>

#include "FilePovGenerator.h"

namespace s2e {
namespace plugins {
namespace pov {

S2E_DEFINE_PLUGIN(FilePovGenerator, "Generates PoVs that fit in one input file", "PovGenerator", "TestCaseGenerator");

void FilePovGenerator::initialize() {
    m_tc = s2e()->getPlugin<testcases::TestCaseGenerator>();

    m_pc = s2e()->getConfig()->getInt(getConfigKey() + ".target_pc", 0x0011223344556677);
    m_gp = s2e()->getConfig()->getInt(getConfigKey() + ".target_gp", 0x8899aabbccddeeff);
}

///
/// \brief Map symbolic decree variables to concrete values.
///
/// The Decree PoV format is designed such that concrete values
/// for registers can be negotiated every time the PoV is executed.
/// Here, we are generating a static file PoV, which cannot have
/// negotiation. We therefore pick hard-coded values for these variables.
///
/// \param varAssignment maps a klee variable name to the actual concrete byte
/// \param opt the PoV generation options
///
void FilePovGenerator::assignNegotiatedVariables(std::unordered_map<std::string, uint8_t> &varAssignment,
                                                 const PovOptions &opt) {
    uint8_t *gp = (uint8_t *) &m_gp;
    uint8_t *pc = (uint8_t *) &m_pc;

    for (auto &it : opt.m_remapping) {
        const auto &kleeVar = it.first;
        const auto &negVar = it.second;

        std::string pc_var = VARNAME_PC + "[";
        std::string gp_var = VARNAME_GP + "[";

        if (boost::starts_with(negVar, gp_var)) {
            unsigned idx = negVar[gp_var.size()] - 0x30;
            if (idx < sizeof(gp)) {
                varAssignment[kleeVar] = gp[idx];
            }
        } else if (boost::starts_with(negVar, pc_var)) {
            unsigned idx = negVar[pc_var.size()] - 0x30;
            if (idx < sizeof(pc)) {
                varAssignment[kleeVar] = pc[idx];
            }
        } else {
            s2e_assert(nullptr, false, "Invalid negotiated variable " << negVar);
        }
    }
}

///
/// \brief Constrain symbolic variables to concrete values for negotiatable registers
///
/// This will force the constraint solver to create a PoV that will deterministically
/// crash the program with the specified concrete values as part of the final register
/// state.
///
/// \param varAssignment the concrete variable assignment
/// \param opt the PoV generation options
///
void FilePovGenerator::constrainNegotiatedVariables(S2EExecutionState *state,
                                                    std::unordered_map<std::string, uint8_t> &varAssignment,
                                                    PovOptions &opt) {
    // Add constraints to force the values for eip and gp register
    for (const auto &it : state->symbolics) {
        auto &array = it;
        if (array->getSize() != 1) {
            getWarningsStream(state) << "Symbolic var " << array->getName() << " must be one byte in size\n";
            continue;
        }

        const auto va = varAssignment.find(array->getName());
        if (va != varAssignment.end()) {
            uint8_t value = (*va).second;
            auto re = klee::ReadExpr::createTempRead(array, klee::Expr::Int8);
            auto cmp = klee::EqExpr::create(re, klee::ConstantExpr::alloc(value, klee::Expr::Int8));
            opt.m_extraConstraints.push_back(cmp);
        }
    }
}

bool FilePovGenerator::generatePoV(S2EExecutionState *state, const PovOptions &_opt, const std::string &filePrefix,
                                   std::vector<std::string> &filePaths) {

    PovOptions opt = _opt;
    std::unordered_map<std::string, uint8_t> varAssignment;

    assignNegotiatedVariables(varAssignment, opt);
    constrainNegotiatedVariables(state, varAssignment, opt);

    auto assignment = klee::Assignment::create();
    if (!solveConstraints(state, opt, *assignment)) {
        getWarningsStream(state) << "Could not solve constraints\n";
        return false;
    }

    testcases::TestCaseData tc;
    m_tc->assembleTestCaseToFiles(*assignment, m_tc->getTemplates(state), tc);

    if (!tc.size()) {
        return false;
    }

    for (const auto &it : tc) {
        const std::string &name = it.first;
        const testcases::Data &data = it.second;

        std::stringstream ss;
        ss << name << "-" << filePrefix << "@" << hexval(_opt.m_faultAddress);

        auto filePath = writeToFile(state, opt, ss.str(), "", data.data(), data.size());
        filePaths.push_back(filePath);
    }

    return true;
}

} // namespace pov
} // namespace plugins
} // namespace s2e
